{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e6c008e-991f-4b45-834e-ded4a5614b4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms\n",
    "from transformers import CLIPTextModel, CLIPTokenizer\n",
    "from diffusers import AutoencoderKL,UNet2DConditionModel\n",
    "from diffusers import PNDMScheduler\n",
    "from tqdm import tqdm\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3dc34778-fe82-4a78-91ef-e4b1f2d0e244",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = AutoencoderKL.from_pretrained(\"/tmp/pretrainmodel/stable-diffusion-v1-5/vae\",)\n",
    "unet = UNet2DConditionModel.from_pretrained(\"/tmp/pretrainmodel/stable-diffusion-v1-5/unet\")\n",
    "tokenizer = CLIPTokenizer.from_pretrained(\"/tmp/pretrainmodel/stable-diffusion-v1-5/tokenizer\")\n",
    "text_encoder = CLIPTextModel.from_pretrained(\"/tmp/pretrainmodel/stable-diffusion-v1-5/text_encoder\")\n",
    "scheduler = PNDMScheduler.from_pretrained(\"/tmp/pretrainmodel/stable-diffusion-v1-5/scheduler\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4cec7194-bae6-4700-8671-2372b555115e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finsh\n"
     ]
    }
   ],
   "source": [
    "device = \"cuda\"\n",
    "vae.to(device)\n",
    "unet.to(device)\n",
    "text_encoder.to(device)\n",
    "print(\"finsh\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "aef70b8d-649c-4ecf-bbf0-a87934f4c7e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 512, 512])\n"
     ]
    }
   ],
   "source": [
    "image = Image.open('images/0.jpg')\n",
    "# 定义转换操作\n",
    "transform = transforms.Compose([\n",
    "    transforms.Resize((512,512)),  # 将图片缩放到224*224\n",
    "    transforms.ToTensor()  # 将PIL Image或numpy.ndarray转换为Tensor，并自动将像素值缩放到[0, 1]并将通道维度提前\n",
    "])\n",
    "# 应用转换\n",
    "image_tensor = transform(image)\n",
    "image_tensor = image_tensor*2-1 # 放缩到(-1,1)\n",
    "image_tensor = image_tensor.unsqueeze(0) \n",
    "# 打印结果\n",
    "print(image_tensor.shape)  # 输出Tensor的形状"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "5fc13574-2a6c-4597-aca9-1c198a4c9b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = [\"A girl's avatar Line art\"] # prompt按自己喜好设置，想生成什么就描述什么\n",
    "height = 512  # SD 默认高\n",
    "width = 512  # SD 默认款\n",
    "num_inference_steps = 20  # 去噪步数\n",
    "guidance_scale = 7.5  # classifier-free guidance (CFG) scale\n",
    "strength = 0.8 # img2img的加噪程度，1则等同于text2img\n",
    "generator=torch.Generator(device=\"cpu\").manual_seed(0) # 为了可复现性，在CPU上生成随机数（注意，要想输出相同的结果，每次都要重置generator）\n",
    "#generator = torch.Generator(device=device)  # 随机种子生成器，用于控制初始的噪声图\n",
    "batch_size = len(prompt)\n",
    "i2i_or_t2i = \"i2i\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "40d1c7f9-baa7-4b99-882a-8e74e4a0d08f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([951, 901, 901, 851, 801, 751, 701, 651, 601, 551, 501, 451, 401, 351,\n",
      "        301, 251, 201, 151, 101,  51,   1], device='cuda:0')\n",
      "final_text_embeddings shape: torch.Size([2, 77, 768])\n",
      "i2i\n",
      "4\n",
      "tensor([801, 751, 701, 651, 601, 551, 501, 451, 401, 351, 301, 251, 201, 151,\n",
      "        101,  51,   1], device='cuda:0')\n",
      "torch.Size([1, 4, 64, 64])\n",
      "tensor(4.2789, device='cuda:0') tensor(-4.2653, device='cuda:0')\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 17/17 [00:03<00:00,  5.15it/s]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "text_input = tokenizer(prompt, padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")\n",
    "max_length = text_input.input_ids.shape[-1]\n",
    "uncond_input = tokenizer([\"\"] * batch_size, padding=\"max_length\", max_length=max_length, return_tensors=\"pt\")\n",
    "scheduler.set_timesteps(num_inference_steps,device=device) # 得到一个list\n",
    "timesteps=scheduler.timesteps\n",
    "print(timesteps)\n",
    "\n",
    "with torch.no_grad():\n",
    "    text_embeddings = text_encoder(text_input.input_ids.to(device))[0]\n",
    "    uncond_embeddings = text_encoder(uncond_input.input_ids.to(device))[0]\n",
    "    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n",
    "    print(\"final_text_embeddings shape:\",text_embeddings.shape)\n",
    "    # print(text_embeddings.sum()) # text_embeddings是一致的\n",
    "    if i2i_or_t2i==\"t2i\":\n",
    "        print(\"t2i\")\n",
    "        latents = torch.randn(\n",
    "            (batch_size, unet.config.in_channels, height // 8, width // 8),\n",
    "            generator=generator,# 为了可复现性，在CPU上创建随机张量\n",
    "            #device=device,\n",
    "        )\n",
    "        latents = latents.to(device)\n",
    "        latents = latents * scheduler.init_noise_sigma # 不同scheduler有不同的系数，对生成效果有影响。i2i不需要？\n",
    "    elif i2i_or_t2i==\"i2i\":\n",
    "        print(\"i2i\")\n",
    "        latents = vae.encode(image_tensor.to(device)).latent_dist.sample()\n",
    "        latents = vae.config.scaling_factor * latents # 0.18215\n",
    "        noise = torch.randn(\n",
    "            latents.shape,\n",
    "            generator=generator,\n",
    "            #device=device,\n",
    "            dtype=text_embeddings.dtype)\n",
    "        noise = noise.to(device)\n",
    "        t_start = int(num_inference_steps-num_inference_steps*strength) # 4\n",
    "        print(t_start)\n",
    "        timesteps = timesteps[t_start:]\n",
    "        print(timesteps)\n",
    "        latents = scheduler.add_noise(latents, noise, timesteps[0])\n",
    "        print(latents.shape)\n",
    "\n",
    "print(latents.max(),latents.min())\n",
    "\n",
    "for t in tqdm(timesteps):\n",
    "    latent_model_input = torch.cat([latents] * 2)\n",
    "\n",
    "    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)\n",
    "\n",
    "    # predict the noise residual\n",
    "    with torch.no_grad():\n",
    "        noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample\n",
    "\n",
    "    # 提示词引导权重cfg\n",
    "    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
    "    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
    "\n",
    "\n",
    "\n",
    "    # compute the previous noisy sample x_t -> x_t-1\n",
    "    latents = scheduler.step(noise_pred, t, latents).prev_sample\n",
    "\n",
    "# 解码图像\n",
    "# scale and decode the image latents with vae\n",
    "latents = 1 / 0.18215 * latents\n",
    "with torch.no_grad():\n",
    "    image = vae.decode(latents).sample\n",
    "\n",
    "image = (image / 2 + 0.5).clamp(0, 1).squeeze()\n",
    "image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()\n",
    "image = Image.fromarray(image)\n",
    "image.save(\"res6.jpg\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5796da8e-2d2e-4137-a01e-f9b5a2f2a8ff",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
